#include <pnm.h>
#include <stdlib.h>
#include <stdio.h>

#include "masked_image.h"
#include "pyramid.h"

/* Algorithm that performs an inpainting function in an input image
 using a masked region from a mask file and saves the result as an 
 output image */

#define PATCH_W 7

/* Clones a pnm image into another image with same dimensions */
int pnm_recopy(pnm* ims, pnm* imd) {

  if(pnm_get_width(*ims) != pnm_get_width(*imd) || 
     pnm_get_height(*ims) != pnm_get_height(*imd))
    return -1;

  for(int i = 0; i < pnm_get_height(*imd); i++) {
    for(int j = 0; j < pnm_get_width(*imd); j++) {

      int r = pnm_get_component(*ims,i,j,PnmRed);
      int g = pnm_get_component(*ims,i,j,PnmGreen);
      int b = pnm_get_component(*ims,i,j,PnmBlue);

      pnm_set_component(*imd,i,j,PnmRed,r);
      pnm_set_component(*imd,i,j,PnmGreen,g);
      pnm_set_component(*imd,i,j,PnmBlue,b);

    }
  }
  return 0;
}

pnm inpainting(char* ims, char* mask, char* imd) {
  
  fprintf(stdout,"Loading initial ppm image: %s\n",ims);
  fprintf(stdout,"Loading mask ppm image: %s\n",mask);

  pnm init = pnm_load(ims);
  pnm immask = pnm_load(mask);
  pnm src = pnm_new(pnm_get_width(init),pnm_get_height(init),PnmRawPpm);
  
  pnm_recopy(&init,&src);

  struct maskedImage* initial = initMaskedImage(&init,&immask);
  struct maskedImage* source = initMaskedImage(&src,&immask);

  struct fifoPyramid imgFifo;

  pushImage(&imgFifo,source);

  /* Building downscaled images pyramid */
  while(pnm_get_width(source->image) > PATCH_W && 
	pnm_get_height(source->image) > PATCH_W) {
    if(countMaskedPixels(source) == 0) break;   
    source = downsample(source);
    pushImage(&imgFifo,source);
  }

  int maxlevel = sizePyramid(imgFifo);

  /* Definining the initial target as the same as the smallest source element
     from the FIFO pyramid. We consider that this target contains no masked 
     pixels */

  /* Cloning the source pnm to link it as a masked Image structure with a void mask */
  pnm copy_source = pnm_new(pnm_get_width(source->image), 
			    pnm_get_height(source->image), PnmRawPpm);

  pnm_recopy(source->image,&copy_source);


  /* Creating a void mask */
  pnm void_mask = pnm_new(pnm_get_width(source->image), 
			  pnm_get_height(source->image), PnmRawPpm);
  
  for(int i = 0; i < pnm_get_height(source->image); i++) {
    for(int j = 0; j < pnm_get_width(source->image); j++) {

      pnm_set_component(void_mask,i,j,PnmRed,0);
      pnm_set_component(void_mask,i,j,PnmGreen,0);
      pnm_set_component(void_mask,i,j,PnmBlue,0);

    }
  }

  /* Linking both pnm */
  struct maskedImage* target = initMaskedImage(&copy_source, &void_mask);

  /* For each level of the pyramid */
  for(int level = maxlevel-1; level>=1; level--) {

    fprintf(stdout,"\n ==== Processing Zoom 1:%d\n",level);

    /* Creating Nearest-Neighbor Fields (direct and reverse) */
    source = popImage(imgFifo);

    fprintf(stdout,"Initialize NNF...");
    
    if(level == maxlevel - 1) {
      
      /* At first, use random data as initial guess */

      // ??? ** TODO ** 

    }

    else {

      /* Then, we use the rebuilt (upscaled) target and reuse the previous
	 NNF as initial guess */

      // ??? ** TODO **
      
    }

    /* Building an upscaled target by ExpectationMaximization-like algorithm */

    // ??? ** TODO ** // target = expectationMaximization(level);

  }

  pnm computed_image = target->image;

  freeMaskedImage(initial);
  freeMaskedImage(target);

  return computed_image;
}

void usage (char *s){
  fprintf(stderr, "Usage: %s <input.ppm> <mask.ppm> <output.ppm> \n", s);
  exit(EXIT_FAILURE);
}

#define ARG_NBR 3

int main(int argc,char** argv) {

  if(argc != ARG_NBR + 1) usage(argv[0]);

  pnm computed_image = inpainting(argv[1], argv[2], argv[3]);

  return EXIT_SUCCESS;
}
